"""用于增强型混合智能体的自适应重规划系统。"""
import re
import time
from typing import Any, Dict, List, Optional, Tuple
from pydantic import BaseModel, Field
from enum import Enum
from json_repair import repair_json
from langchain.schema import HumanMessage

from .planner import Plan, PlanStep, PlanType, Planner
from .executor import ExecutionStatus, StepResult
from .tool_manager import ToolManager
from llm.llm_manager import get_llm_manager
from agent.agent_prompts import replan_prompt_template


class ReplanTrigger(Enum):
    """触发重新规划的原因。"""
    EXECUTION_FAILURE = "execution_failure"
    UNEXPECTED_RESULT = "unexpected_result"
    MISSING_INFORMATION = "missing_information"
    EFFICIENCY_OPTIMIZATION = "efficiency_optimization"
    GOAL_REFINEMENT = "goal_refinement"
    CONTEXT_CHANGE = "context_change"


class AdaptationStrategy(Enum):
    """不同的适应策略。"""
    REPLAN_COMPLETE = "replan_complete"           # 用新计划重新开始
    REPLAN_PARTIAL = "replan_partial"             # 修改剩余步骤
    SWITCH_APPROACH = "switch_approach"           # 在“规划-执行”模式与“反应”模式之间切换（反之亦然）
    ADD_VERIFICATION = "add_verification"         # 添加验证步骤
    PARALLEL_EXECUTION = "parallel_execution"     # 并行尝试多种方法
    INCREMENTAL_SEARCH = "incremental_search"     # 拆分为更小的可搜索单元


class ReplanDecision(BaseModel):
    """关于是否重新规划以及如何重新规划的决策."""
    should_replan: bool
    trigger: Optional[ReplanTrigger]
    strategy: Optional[AdaptationStrategy]
    confidence: float
    reasoning: str
    estimated_improvement: float  # 成功概率的预期提升值
    cost_benefit_ratio: float     # 重新规划的成本与预期收益之比


class AdaptationContext(BaseModel):
    """自适应重新规划决策的上下文"""
    original_query: str
    current_plan: Plan
    execution_results: List[StepResult]
    partial_outputs: Dict[str, Any]
    failed_attempts: List[Dict[str, Any]]
    available_tools: List[str]
    time_budget_remaining: float
    success_probability: float
    context_variables: Dict[str, Any]


class AdaptiveReplanner:
    """基于执行结果反馈进行调整策略的重新规划系统."""
    
    def __init__(self, planner: Planner, tool_manager: ToolManager):
        self.planner = planner
        self.tool_manager = tool_manager
        self.llm_manager = get_llm_manager()
        
        # 用于学习的重新规划历史
        self.replanning_history: List[Dict[str, Any]] = []
        
        # 优化的成功模式
        self.success_patterns: Dict[str, List[Dict[str, Any]]] = {}
        
        # 性能指标
        self.metrics = {
            "total_replans": 0,
            "successful_replans": 0,
            "efficiency_improvements": 0,
            "avg_improvement_ratio": 0.0
        }
    
    async def should_replan(self, context: AdaptationContext, session_id: str = None) -> ReplanDecision:
        """分析执行背景，判断重新规划是否有益."""
        
        # 快速检查是否存在需要重新规划的显著触发条件
        obvious_triggers = self._check_obvious_triggers(context)
        if obvious_triggers:
            return ReplanDecision(
                should_replan=True,
                trigger=obvious_triggers[0],
                strategy=self._recommend_strategy(obvious_triggers[0], context),
                confidence=0.9,
                reasoning=f"Obvious trigger detected: {obvious_triggers[0].value}",
                estimated_improvement=0.7,
                cost_benefit_ratio=3.0
            )
        return await self._analyze_with_llm(context, session_id)
    
    async def execute_adaptive_replan(self, decision: ReplanDecision, context: AdaptationContext, 
                                     session_id: str = None) -> Tuple[Plan, Dict[str, Any]]:
        """执行新规划"""
        start_time = time.time()
        
        try:
            if decision.strategy == AdaptationStrategy.REPLAN_COMPLETE:
                new_plan = await self._replan_complete(context, session_id)
                
            elif decision.strategy == AdaptationStrategy.REPLAN_PARTIAL:
                new_plan = await self._replan_partial(context, session_id)
                
            elif decision.strategy == AdaptationStrategy.SWITCH_APPROACH:
                new_plan = await self._switch_approach(context, session_id)
                
            elif decision.strategy == AdaptationStrategy.ADD_VERIFICATION:
                new_plan = await self._add_verification_steps(context, session_id)
                
            elif decision.strategy == AdaptationStrategy.PARALLEL_EXECUTION:
                new_plan = await self._create_parallel_plan(context, session_id)
                
            elif decision.strategy == AdaptationStrategy.INCREMENTAL_SEARCH:
                new_plan = await self._create_incremental_search_plan(context, session_id)

            else:
                # 降级为部分重新规划
                new_plan = await self._replan_partial(context, session_id)
            
            # 记录重新规划的尝试
            replan_record = {
                "timestamp": time.time(),
                "trigger": decision.trigger.value,
                "strategy": decision.strategy.value,
                "original_plan_id": context.current_plan.id,
                "new_plan_id": new_plan.id,
                "execution_time": time.time() - start_time,
                "confidence": decision.confidence
            }
            
            self.replanning_history.append(replan_record)
            self.metrics["total_replans"] += 1
            
            return new_plan, replan_record
            
        except Exception as e:
            # 若重新规划失败，则返回经修改的原始计划
            fallback_plan = await self._create_fallback_plan(context, str(e), session_id)
            return fallback_plan, {"error": str(e), "fallback": True}
    
    def _check_obvious_triggers(self, context: AdaptationContext) -> List[ReplanTrigger]:
        """检查明显的重新规划触发因素"""
        triggers = []
        
        # 检查执行失败的步骤
        failed_steps = [r for r in context.execution_results if r.status == ExecutionStatus.FAILED]
        if len(failed_steps) >= 2:
            triggers.append(ReplanTrigger.EXECUTION_FAILURE)
        
        # 检查是否缺少关键信息
        if any("not found" in str(r.error).lower() or "missing" in str(r.error).lower() 
               for r in failed_steps):
            triggers.append(ReplanTrigger.MISSING_INFORMATION)
        
        # 检查成功概率
        if context.success_probability < 0.3:
            triggers.append(ReplanTrigger.EFFICIENCY_OPTIMIZATION)
        
        # 检查是否存在过多未取得进展的步骤
        if len(context.execution_results) > len(context.current_plan.steps) * 1.5:
            triggers.append(ReplanTrigger.EFFICIENCY_OPTIMIZATION)
        
        return triggers
    
    def _recommend_strategy(self, trigger: ReplanTrigger, context: AdaptationContext) -> AdaptationStrategy:
        """根据触发条件推荐适配策略."""
        
        if trigger == ReplanTrigger.EXECUTION_FAILURE:
            # 若多个工具均失败，则尝试不同方法
            if len(set(r.step_id for r in context.execution_results if r.status == ExecutionStatus.FAILED)) > 1:
                return AdaptationStrategy.SWITCH_APPROACH
            else:
                return AdaptationStrategy.REPLAN_PARTIAL
        
        elif trigger == ReplanTrigger.MISSING_INFORMATION:
            return AdaptationStrategy.INCREMENTAL_SEARCH
        
        elif trigger == ReplanTrigger.EFFICIENCY_OPTIMIZATION:
            return AdaptationStrategy.PARALLEL_EXECUTION
        
        elif trigger == ReplanTrigger.UNEXPECTED_RESULT:
            return AdaptationStrategy.ADD_VERIFICATION
        
        else:
            return AdaptationStrategy.REPLAN_PARTIAL
    
    async def _analyze_with_llm(self, context: AdaptationContext, session_id: str = None) -> ReplanDecision:
        """使用大语言模型分析重新规划是否有益"""

        prompt = self._create_analysis_prompt(context)
        messages = [
            HumanMessage(content=prompt)
        ]
        
        try:
            response = await self.llm_manager.ainvoke(messages, session_id)
            analysis_text = response.content

            return self._parse_replan_decision(analysis_text)
            
        except Exception as e:
            # 保守性备用方案：除非存在明显问题，否则不重新规划
            return ReplanDecision(
                should_replan=len([r for r in context.execution_results if r.status == ExecutionStatus.FAILED]) >= 2,
                trigger=ReplanTrigger.EXECUTION_FAILURE if context.execution_results else None,
                strategy=AdaptationStrategy.REPLAN_PARTIAL,
                confidence=0.5,
                reasoning=f"LLM 分析失败：{str(e)}，采用保守方法。",
                estimated_improvement=0.3,
                cost_benefit_ratio=1.0
            )
    
    async def _replan_complete(self, context: AdaptationContext, session_id: str = None) -> Plan:
        """从头开始制定一份全新的方案"""
        
        # 分析原方案存在的问题
        failure_analysis = self._analyze_failures(context.execution_results)
        
        # 结合经验教训制定新方案
        enhanced_context = {
            **context.context_variables,
            "previous_failures": failure_analysis,
            "avoid_tools": [r.step_id for r in context.execution_results if r.status == ExecutionStatus.FAILED],
            "successful_partial_results": context.partial_outputs
        }
        
        return await self.planner.create_plan(
            query=context.original_query,
            available_tools=context.available_tools,
            context=enhanced_context,
            session_id=session_id
        )
    
    async def _replan_partial(self, context: AdaptationContext, session_id: str = None) -> Plan:
        """修改当前方案的剩余步骤"""
        
        # 识别已完成的步骤
        completed_steps = [r.step_id for r in context.execution_results if r.status == ExecutionStatus.COMPLETED]
        
        # 使用现有规划器的 refine_plan 方法制定优化方案
        execution_results_dict = [
            {
                "step_id": r.step_id,
                "success": r.status == ExecutionStatus.COMPLETED,
                "error": r.error or "",
                "output": r.output
            }
            for r in context.execution_results
        ]
        
        return await self.planner.refine_plan(
            plan=context.current_plan,
            execution_results=execution_results_dict,
            current_context=context.context_variables,
            session_id=session_id
        )
    
    async def _switch_approach(self, context: AdaptationContext, session_id: str = None) -> Plan:
        """结构化规划 <-> 响应式方法"""
        
        # 制定一份采用更动态的工具调用方式的简易响应式方案
        steps = []
        
        # 添加动态搜索 / 探索步骤
        steps.append(PlanStep(
            id="dynamic_search",
            description="使用工具动态搜索信息",
            tool="wikipedia",
            input_template=context.original_query,
            dependencies=[],
            confidence=0.7
        ))
        
        # 若查询涉及数值，则添加计算步骤
        if any(word in context.original_query.lower() for word in ["calculate", "compute", "number", "math"]):
            steps.append(PlanStep(
                id="dynamic_calc",
                description="执行所有必要的计算",
                tool="calculator",
                input_template="根据之前的结果进行计算：{{search_result}}",
                dependencies=["dynamic_search"],
                confidence=0.6
            ))
        
        return Plan(
            id=f"switched_plan_{int(time.time())}",
            query=context.original_query,
            goal="采用渐进式响应式方法动态解决查询问题",
            plan_type=PlanType.SEQUENTIAL,
            steps=steps,
            estimated_duration=45.0,
            confidence=0.6,
            metadata={"approach": "switched", "original_plan": context.current_plan.id},
            created_at=time.time()
        )
    
    async def _add_verification_steps(self, context: AdaptationContext, session_id: str = None) -> Plan:
        """添加验证和确认步骤以提高可靠性。"""
        
        # 复制当前方案并添加验证步骤
        import copy
        new_plan = copy.deepcopy(context.current_plan)
        new_plan.id = f"verified_plan_{int(time.time())}"
        
        # 在每个主要步骤后添加验证步骤
        verification_steps = []
        for step in new_plan.steps:
            if step.tool in ["calculator", "database"]:  # Add verification for critical tools
                verify_step = PlanStep(
                    id=f"verify_{step.id}",
                    description=f"Verify the result from {step.id}",
                    tool="calculator" if step.tool == "calculator" else "database",
                    input_template=f"Verify: {{output_from_{step.id}}}",
                    dependencies=[step.id],
                    confidence=0.8
                )
                verification_steps.append(verify_step)
        
        new_plan.steps.extend(verification_steps)
        new_plan.metadata["verification_added"] = True
        
        return new_plan
    
    async def _create_parallel_plan(self, context: AdaptationContext, session_id: str = None) -> Plan:
        """制定一份多路并行执行的计划"""
        
        steps = []
        
        # 制定信息收集的并行方案
        steps.append(PlanStep(
            id="parallel_search_1",
            description="搜索方法 1：一般性信息",
            tool="wikipedia",
            input_template=context.original_query,
            dependencies=[],
            confidence=0.7
        ))
        
        steps.append(PlanStep(
            id="parallel_search_2", 
            description="搜索方法 2：网络搜索",
            tool="web_search",
            input_template=context.original_query,
            dependencies=[],
            confidence=0.7
        ))

        steps.append(PlanStep(
            id="synthesize_results",
            description="整合并行搜索的结果",
            tool="calculator",  # 使用计算器处理 / 合并结果
            input_template="Combine: {{parallel_search_1}} and {{parallel_search_2}}",
            dependencies=["parallel_search_1", "parallel_search_2"],
            confidence=0.8
        ))
        
        return Plan(
            id=f"parallel_plan_{int(time.time())}",
            query=context.original_query,
            goal="通过并行信息收集来解决查询问题",
            plan_type=PlanType.PARALLEL,
            steps=steps,
            estimated_duration=30.0,
            confidence=0.8,
            metadata={"parallel_execution": True},
            created_at=time.time()
        )
    
    async def _create_incremental_search_plan(self, context: AdaptationContext, session_id: str = None) -> Plan:
        """制定一份将查询拆分为渐进式可搜索部分的方案"""
        
        # 将查询拆分为更小的、可搜索的组成部分
        query_parts = self._decompose_query(context.original_query)
        
        steps = []
        for i, part in enumerate(query_parts):
            steps.append(PlanStep(
                id=f"incremental_search_{i+1}",
                description=f"Search for: {part}",
                tool="wikipedia",
                input_template=part,
                dependencies=[f"incremental_search_{j+1}" for j in range(i)],
                confidence=0.8
            ))
        
        # 最终综合步骤
        steps.append(PlanStep(
            id="synthesize_incremental",
            description="整合所有渐进式搜索结果",
            tool="calculator",
            input_template="整合所有搜索结果用于回答： " + context.original_query,
            dependencies=[f"incremental_search_{i+1}" for i in range(len(query_parts))],
            confidence=0.9
        ))
        
        return Plan(
            id=f"incremental_plan_{int(time.time())}",
            query=context.original_query,
            goal="通过渐进式搜索与综合来解决查询问题",
            plan_type=PlanType.SEQUENTIAL,
            steps=steps,
            estimated_duration=60.0,
            confidence=0.8,
            metadata={"incremental_search": True, "query_parts": query_parts},
            created_at=time.time()
        )
    
    def _decompose_query(self, query: str) -> List[str]:
        """将复杂查询分解为可搜索的部分"""
        
        # 简单启发式分解（可通过自然语言处理（NLP）进行增强）
        parts = []
        if re.search("( and |并且)", query, re.IGNORECASE) is not None:
            parts.extend(re.split("( and |并且)", query, re.IGNORECASE))
        elif ", " in query:
            parts.extend(query.split(", "))
        elif re.search("(then|那么)", query, re.IGNORECASE) is not None:
            parts.extend(re.split("(then|那么)", query, re.IGNORECASE))
        else:
            # 尝试识别关键概念
            words = query.split()
            # 若查询内容较长，则拆分为多个片段，否则不拆分
            if len(words) > 6:
                chunk_size = max(3, len(words) // 3)
                for i in range(0, len(words), chunk_size):
                    parts.append(" ".join(words[i:i+chunk_size]))
            else:
                parts = [query]
        
        return [part.strip() for part in parts if part.strip()]
    
    def _analyze_failures(self, execution_results: List[StepResult]) -> Dict[str, Any]:
        """分析执行结果的失败原因."""
        failures = [r for r in execution_results if r.status == ExecutionStatus.FAILED]
        
        analysis = {
            "total_failures": len(failures),
            "failed_tools": list(set(f.step_id for f in failures)),
            "common_errors": {},
            "failure_patterns": []
        }

        for failure in failures:
            error_key = str(failure.error).lower()[:50]
            analysis["common_errors"][error_key] = analysis["common_errors"].get(error_key, 0) + 1

        if len(analysis["failed_tools"]) == 1:
            analysis["failure_patterns"].append("single_tool_repeated_failure")
        elif analysis["total_failures"] > len(execution_results) / 2:
            analysis["failure_patterns"].append("high_failure_rate")
        
        return analysis
    
    def _create_analysis_prompt(self, context: AdaptationContext) -> str:
        """创建用于大语言模型（LLM）重规划决策的提示词"""
        
        results_summary = []
        for result in context.execution_results:
            status = "SUCCESS" if result.status == ExecutionStatus.COMPLETED else "FAILED"
            results_summary.append(f"Step {result.step_id}: {status} - {result.error or 'OK'}")

        return replan_prompt_template.format(
            inputs={
                'original_query': context.original_query,
                'goal': context.current_plan.goal,
                'plan_type': context.current_plan.plan_type.value,
                'steps_count': len(context.current_plan.steps),
                'results_summary': '\n'.join(results_summary),
                'success_probability': f'{context.success_probability:.2f}',
                'time_budget_remaining': f'{context.time_budget_remaining:.1f}'
            },
            remove_template_variables=True
        )

    def _parse_replan_decision(self, analysis_text: str) -> ReplanDecision:
        """将大语言模型分析结果解析为重新规划决策"""
        
        try:
            import re
            json_match = re.search(r'\{.*\}', analysis_text, re.DOTALL)
            if json_match:
                data = repair_json(json_match.group(), return_objects=True)
                
                return ReplanDecision(
                    should_replan=data.get("should_replan", False),
                    trigger=ReplanTrigger(data.get("trigger", "execution_failure")) if data.get("trigger") else None,
                    strategy=AdaptationStrategy(data.get("strategy", "replan_partial")) if data.get("strategy") else None,
                    confidence=float(data.get("confidence", 0.5)),
                    reasoning=data.get("reasoning", "LLM analysis"),
                    estimated_improvement=float(data.get("estimated_improvement", 0.3)),
                    cost_benefit_ratio=float(data.get("cost_benefit_ratio", 2.0))
                )
        except Exception:
            pass

        return ReplanDecision(
            should_replan=False,
            trigger=None,
            strategy=None,  
            confidence=0.3,
            reasoning="大语言模型解析失败",
            estimated_improvement=0.0,
            cost_benefit_ratio=0.5
        )

    async def _create_fallback_plan(self, context: AdaptationContext, error: str, session_id: str = None) -> Plan:
        """重新规划失败时，制定一份简单的备用方案"""

        return Plan(
            id=f"fallback_plan_{int(time.time())}",
            query=context.original_query,
            goal="Simple fallback approach",
            plan_type=PlanType.SEQUENTIAL,
            steps=[
                PlanStep(
                    id="fallback_step",
                    description="采用简单方法处理查询",
                    tool=context.available_tools[0] if context.available_tools else "calculator",
                    input_template=context.original_query,
                    dependencies=[],
                    confidence=0.4
                )
            ],
            estimated_duration=20.0,
            confidence=0.4,
            metadata={"fallback": True, "replan_error": error},
            created_at=time.time()
        )
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取重新规划的性能指标"""
        success_rate = (self.metrics["successful_replans"] / max(1, self.metrics["total_replans"]))
        
        return {
            **self.metrics,
            "success_rate": success_rate,
            "recent_replans": self.replanning_history[-10:] if self.replanning_history else []
        }